from grafanalib.core import AlertGroup
from grafanalib._gen import DashboardEncoder, loader
import json
import requests
from os import getenv


def get_alert_json(alert: AlertGroup):
    '''
    get_alert_json generates JSON from grafanalib AlertGroup object

    :param alert - AlertGroup created via grafanalib
    '''

    return json.dumps(alert.to_json_data(), sort_keys=True, indent=4, cls=DashboardEncoder)


def upload_to_grafana(alertjson, folder, server, api_key, session_cookie, verify=True):
    '''
    upload_to_grafana tries to upload the AlertGroup to grafana and prints response
    WARNING: This will first delete all alerts in the AlertGroup before replacing them with the provided AlertGroup.

    :param alertjson - AlertGroup json generated by grafanalib
    :param folder - Folder to upload the AlertGroup into
    :param server - grafana server name
    :param api_key - grafana api key with read and write privileges
    '''
    groupName = json.loads(alertjson)['name']

    headers = {}
    if api_key:
        print("using bearer auth")
        headers['Authorization'] = f"Bearer {api_key}"

    if session_cookie:
        print("using session cookie")
        headers['Cookie'] = session_cookie

    print(f"deleting AlertGroup {groupName} in folder {folder}")
    r = requests.delete(f"https://{server}/api/ruler/grafana/api/v1/rules/{folder}/{groupName}", headers=headers, verify=verify)
    print(f"{r.status_code} - {r.content}")

    headers['Content-Type'] = 'application/json'

    print(f"ensuring folder {folder} exists")
    r = requests.post(f"https://{server}/api/folders", data={"title": folder}, headers=headers)
    print(f"{r.status_code} - {r.content}")

    print(f"uploading AlertGroup {groupName} to folder {folder}")
    r = requests.post(f"https://{server}/api/ruler/grafana/api/v1/rules/{folder}", data=alertjson, headers=headers, verify=verify)
    # TODO: add error handling
    print(f"{r.status_code} - {r.content}")


grafana_api_key = getenv("GRAFANA_API_KEY")
grafana_server = getenv("GRAFANA_SERVER")
grafana_cookie = getenv("GRAFANA_COOKIE")

# Generate an alert from the example
my_alergroup_json = get_alert_json(loader("./grafanalib/tests/examples/example.alertgroup.py"))
upload_to_grafana(my_alergroup_json, "testfolder", grafana_server, grafana_api_key, grafana_cookie)
